#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project: AI_Testing
@File: ai_utils.py
@Author: Safir
@Date: 2025/3/20 10:54
'''
import json
import requests
import time
import logging
from backend.app.config.settings import AI_MODEL_CONFIG

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class AIModelClient:
    """AI模型客户端，用于与不同的AI模型API交互"""

    @staticmethod
    def get_client(model_type=None):
        """获取指定类型的AI模型客户端"""
        if model_type is None:
            model_type = AI_MODEL_CONFIG.get("default_model", "deepseek")
        if model_type == "deepseek":
            return DeepseekClient()
        elif model_type == "douban":
            return DoubanClient()
        else:
            raise ValueError(f"不支持的模型类型: {model_type}")


class DeepseekClient:
    """DeepSeek API客户端"""

    def __init__(self):
        self.api_key = AI_MODEL_CONFIG["deepseek"]["api_key"]
        self.model_name = AI_MODEL_CONFIG["deepseek"]["model_name"]
        self.api_url = "https://api.deepseek.com/v1/chat/completions"

    def generate(self, prompt, max_tokens=8192, temperature=0.7, top_p=0.8):
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

        payload = {
            "model": self.model_name,
            "messages": [
                {"role": "system",
                 "content": "你是一个专业的网络空间靶场系统软件的 测试和需求分析专家，擅长在简单的需求的基础上优化需求和根据完整的需求生成完整的测试用例。"},
                {"role": "user", "content": prompt}
            ],
            "max_tokens": max_tokens,
            "temperature": temperature,
            "top_p": top_p
        }

        max_retries = 3
        retries = 0
        while retries < max_retries:
            try:
                start_time = time.time()
                response = requests.post(self.api_url, headers=headers, json=payload, timeout=100)
                response.raise_for_status()
                result = response.json()

                logger.info(f"DeepSeek API调用耗时: {time.time() - start_time:.2f}秒")

                # 解析响应
                if "choices" in result and len(result["choices"]) > 0:
                    return result["choices"][0]["message"]["content"]
                else:
                    logger.error(f"DeepSeek API返回异常结果: {result}")
                    return "无法从API获取有效响应。"
            except requests.exceptions.Timeout:
                logger.error("DeepSeek API调用超时，正在重试...")
                retries += 1
                time.sleep(2)  # 等待 2 秒后重试
            except requests.exceptions.RequestException as e:
                logger.error(f"DeepSeek API请求异常: {str(e)}")
                break
            except Exception as e:
                logger.error(f"DeepSeek API调用失败: {str(e)}")
                break

        return "无法从API获取有效响应。"


class DoubanClient:
    """豆包(DashScope) API客户端"""

    def __init__(self):
        self.api_key = AI_MODEL_CONFIG["douban"]["api_key"]
        self.model_name = AI_MODEL_CONFIG["douban"]["model_name"]
        self.api_url = "https://api.dashscope.com/v1/chat/completions"

    def generate(self, prompt, max_tokens=4096, temperature=0.7, top_p=0.8):
        """生成文本"""
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

        payload = {
            "model": self.model_name,
            "messages": [
                {"role": "system",
                 "content": "你是一个专业的网络空间靶场系统软件的 测试和需求分析专家，擅长在简单的需求的基础上优化需求和根据完整的需求生成完整的测试用例。"},
                {"role": "user", "content": prompt}
            ],
            "parameters": {
                "max_tokens": max_tokens,
                "temperature": temperature,
                "top_p": top_p
            }
        }

        try:
            start_time = time.time()
            response = requests.post(self.api_url, headers=headers, json=payload, timeout=30)
            response.raise_for_status()
            result = response.json()

            logger.info(f"豆包API调用耗时: {time.time() - start_time:.2f}秒")

            # 解析响应
            if "output" in result and "choices" in result["output"] and len(result["output"]["choices"]) > 0:
                return result["output"]["choices"][0]["message"]["content"]
            else:
                logger.error(f"豆包API返回异常结果: {result}")
                return "无法从API获取有效响应。"
        except requests.exceptions.Timeout:
            logger.error("豆包API调用超时")
            return "API调用超时，请稍后重试。"
        except requests.exceptions.RequestException as e:
            logger.error(f"豆包API请求异常: {str(e)}")
            return f"API请求异常: {str(e)}"
        except Exception as e:
            logger.error(f"豆包API调用失败: {str(e)}")
            return f"API调用失败: {str(e)}"
